#include "kernel_operator.h"
#include "nll_loss_1d.h"
#include "nll_loss_sca.h"
#include "nll_loss_new.h"
#include "nll_loss_sum.h"

using namespace AscendC;

extern "C" __global__ __aicore__ void nll_loss(GM_ADDR x, GM_ADDR target,
                                               GM_ADDR weight, GM_ADDR y,
                                               GM_ADDR workspace,
                                               GM_ADDR tiling)
{
  GET_TILING_DATA(tiling_data, tiling);
  // TODO: user kernel impl
  TPipe pipe;
  KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); // 增加这一行
  // printf("-----------------------Processing running\n");
  // printf("dim=%d N,C,M = {%d, %d, %d} \n reduction = %d ignore = %d \n",
  //        tiling_data.dimNum, tiling_data.sizeN, tiling_data.sizeC,
  //        tiling_data.sizeM, tiling_data.reduction, tiling_data.ignore_index);
  // printf(" coreIterNum, IterModCore, coreNum = {%d %d %d}\n",
  //        tiling_data.coreIterNum, tiling_data.IterModCore,
  //        tiling_data.coreNum);

  if (TILING_KEY_IS(1)) // 一维数据
  {
    if (GetBlockIdx() != 0)
      return;
    NLLLoss1d op1;
    op1.Init(x, target, weight, y, tiling_data.sizeC,
             tiling_data.reduction, tiling_data.ignore_index);
    op1.process();
  }
  else if (TILING_KEY_IS(2)) // mean case
  {
    // NLLLossNew op;
    // op.Init(x, target, weight, y, workspace, tiling_data.reduction,
    //         tiling_data.ignore_index,
    //         tiling_data.sizeC, tiling_data.coreRowNumSmall, tiling_data.coreRowNumBig, tiling_data.RowModCore,
    //         tiling_data.coreNum);
    // op.process();
    if (GetBlockIdx() != 0)
      return;
    NLLLossSca opsca;
    opsca.Init(x, target, weight, y, tiling_data.reduction,
               tiling_data.ignore_index, tiling_data.sizeN, tiling_data.sizeC,
               1);
    opsca.process();
  }
  else if (TILING_KEY_IS(3)) // sum case
  {
    NLLLossSum op;
    op.Init(x, target, weight, y,tiling_data.ignore_index, tiling_data.coreNum, tiling_data.sizeC, tiling_data.coreRowNum , tiling_data.coreRowMod,&pipe);
    op.process();
  }
  else if (TILING_KEY_IS(4))
  {
    printf("dim error! dim = %d\n", tiling_data.dimNum);
  }

  // if (tiling_data.dimNum == 2)
  // {
  //   //     NLLLoss op;
  //   //     op.Init(x, target, weight, y, workspace, tiling_data.reduction,
  //   //             tiling_data.ignore_index, tiling_data.dimNum, tiling_data.sizeN,
  //   //             tiling_data.sizeC, tiling_data.coreIterNum, tiling_data.IterModCore,
  //   //             tiling_data.coreNum);
  //   //     op.process();

  //   //     op.processFinal();
  //   if (tiling_data.reduction == 1)
  //   {
  //     NLLLossSum op;
  //     op.Init(x, target, weight, y, workspace,
  //             tiling_data.ignore_index,
  //             tiling_data.sizeC, tiling_data.coreRowNumSmall, tiling_data.coreRowNumBig, tiling_data.RowModCore,
  //             tiling_data.coreNum);
  //     op.process();
  //   }
  //   else
  //   {
  //     NLLLossNew op;
  //     op.Init(x, target, weight, y, workspace, tiling_data.reduction,
  //             tiling_data.ignore_index,
  //             tiling_data.sizeC, tiling_data.coreRowNumSmall, tiling_data.coreRowNumBig, tiling_data.RowModCore,
  //             tiling_data.coreNum);
  //     op.process();
  //   }

  //   // op.processtest();
  // }
  // else if (tiling_data.dimNum == 1)
  // {
  //   if (GetBlockIdx() != 0)
  //     return;
  //   NLLLoss1d op1;
  //   op1.Init(x, target, weight, y, tiling_data.sizeC,
  //            tiling_data.reduction, tiling_data.ignore_index);
  //   op1.process();
  // }
  // else
  // {
  //   if (GetBlockIdx() != 0)
  //     return;
  //   NLLLossSca opsca;

  //   opsca.Init(x, target, weight, y, tiling_data.reduction,
  //              tiling_data.ignore_index, tiling_data.sizeN, tiling_data.sizeC,
  //              tiling_data.sizeM);
  //   opsca.process();
  // }
  // printf("--------------down------------\n");
  return;
}